package com.example.scheduler;

import java.util.*;

/**
 * DagTemplate：描述一个DAG“模板”及其依赖关系。
 * 通过内置的邻接表和入度表，可使用拓扑排序将DAG拆分为若干层级。
 */
public class DagTemplate {
    private final String templateId;
    // adjacencyList：每个任务 -> 其后继任务列表
    private final Map<String, List<String>> adjacencyList;
    // inDegree：每个任务的依赖数（有多少前置任务）
    private final Map<String, Integer> inDegree;

    public DagTemplate(String templateId) {
        this.templateId = templateId;
        this.adjacencyList = new HashMap<>();
        this.inDegree = new HashMap<>();
    }

    public String getTemplateId() {
        return templateId;
    }

    /**
     * 添加一个独立的任务（若不存在）
     */
    public void addTask(String taskId) {
        adjacencyList.putIfAbsent(taskId, new ArrayList<>());
        inDegree.putIfAbsent(taskId, 0);
    }

    /**
     * 添加依赖：prerequisite -> task
     * 表示 "task" 需要在 "prerequisite" 完成之后才能执行
     */
    public void addDependency(String prerequisite, String task) {
        addTask(prerequisite);
        addTask(task);
        adjacencyList.get(prerequisite).add(task);
        inDegree.put(task, inDegree.get(task) + 1);
    }

    /**
     * 使用Kahn算法对DAG进行拓扑排序，并按层级划分
     * 返回：List<List<String>>，其中每个子List代表一层可并发执行的任务
     */
    public List<List<String>> getLevels() throws Exception {
        // 拷贝 inDegree，用于过程处理
        Map<String, Integer> inDeg = new HashMap<>(inDegree);
        List<List<String>> levels = new ArrayList<>();
        Queue<String> queue = new LinkedList<>();

        // 将入度为0的节点（无依赖任务）入队
        for (Map.Entry<String, Integer> e : inDeg.entrySet()) {
            if (e.getValue() == 0) {
                queue.offer(e.getKey());
            }
        }
        if (queue.isEmpty()) {
            throw new Exception("No start node found in template " + templateId + " (possible cycle)");
        }

        // 按层遍历
        while (!queue.isEmpty()) {
            int size = queue.size();
            List<String> currentLevel = new ArrayList<>();
            for (int i = 0; i < size; i++) {
                String task = queue.poll();
                currentLevel.add(task);
                // 遍历该task的所有后继节点
                for (String neighbor : adjacencyList.getOrDefault(task, Collections.emptyList())) {
                    inDeg.put(neighbor, inDeg.get(neighbor) - 1);
                    if (inDeg.get(neighbor) == 0) {
                        queue.offer(neighbor);
                    }
                }
            }
            levels.add(currentLevel);
        }

        // 若有节点未处理完，说明图中有环
        int processedCount = levels.stream().mapToInt(List::size).sum();
        if (processedCount != adjacencyList.size()) {
            throw new Exception("Cycle detected in template " + templateId);
        }

        return levels;
    }
}